import tiktoken
import json
from tqdm import tqdm
encoding = tiktoken.get_encoding("cl100k_base")
def calc_token(text):
    tokens = encoding.encode(text)
    return len(tokens)
def merge(jsonl_file_path, output_file_path=None):
    """
    读取JSONL文件，将所有对话内容合并成一个文本
    
    Args:
        jsonl_file_path: 输入的JSONL文件路径
        output_file_path: 输出的文本文件路径（可选，如果不提供则只返回文本内容）
    
    Returns:
        merged_text: 合并后的文本内容
    """
    txts=[]
    try:
        with open(jsonl_file_path, 'r', encoding='utf-8') as file:
            for line_num, line in enumerate(file, 1):
                line = line.strip()
                if not line:
                    continue
                    
                try:
                    data = json.loads(line)
                    messages = data.get('messages', [])
                    merged_text=''
                    # 处理每条消息
                    for i, message in enumerate(messages):
                        role = message.get('role', 'unknown')
                        content = message.get('content', '')
                        
                        # 添加消息到合并文本
                        merged_text += f"{content}"
                    
                    txts.append(merged_text)
                    
                except json.JSONDecodeError as e:
                    print(f"第{line_num}行JSON解析错误: {e}")
                except Exception as e:
                    print(f"第{line_num}行处理错误: {e}")
        
        # 如果提供了输出文件路径，则保存到文件
        if output_file_path:
            with open(output_file_path, 'w', encoding='utf-8') as f:
                f.write(merged_text)
            print(f"合并后的文本已保存到: {output_file_path}")
        
        return txts
        
    except Exception as e:
        print(f"读取文件时发生错误: {e}")
        return ""
path=input("输入jsonl路径：")
alltxt=merge(path)
cnt=0
for txt in tqdm(alltxt):
    cnt+=calc_token(txt)
print(f'根据OpenAI的tokenizer，token总量为{cnt}，如果要微调gpt-4o-2024-08-06的话，价格是{25*cnt/1000000}$，如果要微调gpt-4o-mini-2024-07-18的话，价格是{3*cnt/1000000}$')